# Support segmentation for multiple languages such as Python, JavaScript, Java, C, GO, etc

import re
import os

def get_language(input_file):
    # 根据文件扩展名确定语言
    _, ext = os.path.splitext(input_file)
    if ext == '.py':
        return 'Python'
    elif ext in ('.js', '.jsx'):
        return 'JavaScript'
    elif ext == '.java':
        return 'Java'
    elif ext in ('.c', '.cpp', '.cc', '.h', '.hpp'):
        return 'C/C++'
    elif ext == '.go':
        return 'Go'
    else:
        raise ValueError(f"Unsupported file type: {ext}")

def segment_code(code, language, max_lines=100):
    segments = []
    current_segment = []
    function_started = False

    if language == 'Python':
        def is_function_start(line):
            return line.strip().startswith("def ")
        
        def is_class_start(line):
            return line.strip().startswith("class ")
    
    elif language in ('JavaScript', 'Java', 'C/C++'):
        def is_function_start(line):
            # 检查是否是函数定义的开始
            return re.match(r'\s*(public|private|protected|\w+)?\s*\w+\s+\w+\s*\(.*?\)\s*\{', line) or \
                   re.match(r'\s*\w+\s+\w+\s*\(.*?\)\s*\{', line)
        
        def is_class_start(line):
            # 检查是否是类定义的开始
            return re.match(r'\s*(public|private|protected)?\s*class\s+\w+', line) or \
                   re.match(r'\s*struct\s+\w+', line)
        
    elif language == 'Go':
        def is_function_start(line):
            # 检查是否是函数定义的开始
            return re.match(r'\s*func\s+\w+\s*\(.*?\)\s*{', line)

        def is_class_start(line):
            # Go 语言没有类，只有结构体
            return re.match(r'\s?type\s+\w+\s+struct', line)
    
    else:
        raise ValueError(f"Unsupported language: {language}")
    
    for line in code.split('\n'):
        if (is_function_start(line) or is_class_start(line)) and current_segment:
            # 如果当前段落不为空且遇到了新的函数或类定义，保存当前段落
            segments.append('\n'.join(current_segment))
            current_segment = []
        
        current_segment.append(line)
        
        if is_function_start(line) or is_class_start(line):
            function_started = True

        if len(current_segment) >= max_lines:
            # 如果当前段落达到最大行数，保存当前段落
            segments.append('\n'.join(current_segment))
            current_segment = []
            function_started = False

    # 添加最后一个段落（如果存在）
    if current_segment:
        segments.append('\n'.join(current_segment))
    
    return segments

def create_output_directory(output_prefix):
    """检查并创建以 output_prefix 命名的目录"""
    output_dir = os.path.dirname(output_prefix)
    if not output_dir:
        # 如果没有指定目录，则使用当前目录
        output_dir = '.'
    
    # 创建以 output_prefix 命名的目录
    final_output_dir = os.path.join(output_dir, os.path.basename(output_prefix))
    if not os.path.exists(final_output_dir):
        os.makedirs(final_output_dir)
        print(f"Created output directory: {final_output_dir}")
    
    return final_output_dir

def main(input_file, output_prefix):
    if not os.path.isfile(input_file):
        raise FileNotFoundError(f"Input file does not exist: {input_file}")
    
    with open(input_file, 'r', encoding='utf-8') as file:
        code = file.read()
    
    language = get_language(input_file)
    segments = segment_code(code, language)

    final_output_dir = create_output_directory(output_prefix)
    
    for i, segment in enumerate(segments):
        # 确保输出文件路径在指定的目录中
        output_filename = os.path.join(final_output_dir, f"{os.path.basename(output_prefix)}_{i}.txt")
        with open(output_filename, 'w', encoding='utf-8') as out_file:
            out_file.write(segment)

if __name__ == "__main__":
    import sys

    if len(sys.argv) != 3:
        print("Usage: python code_segmentation.py <input_file> <output_prefix>")
        sys.exit(1)
    
    input_file = sys.argv[1]
    output_prefix = sys.argv[2]
    
    main(input_file, output_prefix)